#include <sds/sds.h>

#include <cluster.h>
#include <matrix.h>

#include <vector>
#include <iostream>
#include <iomanip>

using namespace sds;
using namespace std;

int main(int argc, const char* argv[]) {
  SFile reffile(argv[2]);
  SDataView<unsigned char> ref =
    reffile.getDataset(0).getDataView<unsigned char>();

  vector<SInteger> size = ref.getSize();

  SFile srcfile(argv[1]);
  SDataView<Point> src(size);

  for (int k = 0; k < CH; ++k) {
    SDataView<unsigned char> view =
      srcfile.getDataset(k).getDataView<unsigned char>();
    for (int i = 0; i < size[0]; ++i) {
      for (int j = 0; j < size[1]; ++j) {
	((Point&)src[i][j])[k] = (double)view[i][j];
      }    
    }
  }

  vector<Cluster> clusters(256);
  for (int i = 0; i < size[0]; ++i) {
    for (int j = 0; j < size[1]; ++j) {
      if (ref[i][j] != 0) {
	clusters[ref[i][j]] += src[i][j];
      }
    }
  }

  vector<int> index;
  vector<Vector> expected;
  vector<Matrix> inverse;

  for (int k = 1; k < 256; ++k) {
    if (clusters[k].size() == 0) continue;
    index.push_back(k);
    Vector exp(CH);
    for (int i = 0; i < CH; ++i) {
      exp(i) = clusters[k].expected(i);
    }
    expected.push_back(exp);
    Matrix cov(CH, CH);
    for (int i = 0; i < CH; ++i) {
      for (int j = 0; j < CH; ++j) {
	cov(i, j) = clusters[k].covariance(i, j);
      }
    }
    inverse.push_back(cov.inverse());
  }

  SDataView<unsigned char> trg(size);

  for (int i = 0; i < size[0]; ++i) {
    for (int j = 0; j < size[1]; ++j) {
      Vector current(CH);
      for (int l = 0; l < CH; ++l) {
	current(l) = ((Point&)src[i][j])[l];
      }
      int minPos = -1; 
      double minVal = 0.0;
      for (int k = 0; k < index.size(); ++k) {
	Vector diff = current - expected[k];
	double dist = diff * inverse[k] * diff;
	if (minPos == -1 || minVal > dist) {
	  minPos = k;
	  minVal = dist;
	}
      }
      if (minVal < 64.0) {
	trg[i][j] = index[minPos];
      } else {
	trg[i][j] = 0;
      }
    }
  }
  
  SFile trgfile(argv[3], SFile::CREATE);
  trgfile.createDataset("clasified", UINT8, size).
    setDataView<unsigned char>(trg);
  return 0;
}
